package com.lxq.scmserver.controller;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.lxq.scmserver.service.BookingTools;
import com.lxq.scmserver.service.LoggingAdvisor;
import java.time.LocalDate;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@RestController
@CrossOrigin
@RequestMapping("/ai")
@Slf4j
public class OpenAiController {

    private final ChatClient dashScopeChatModel;
    private final ChatClient ollamaChatModel;


    public OpenAiController(DashScopeChatModel dashScopeChatModel,
                            OllamaChatModel ollamaChatModel,
                            ToolCallbackProvider mcpTools,
                            BookingTools bookingTools,
                            ChatMemory chatMemory) {

        this.dashScopeChatModel = ChatClient.builder(dashScopeChatModel)
                .defaultSystem("""
                        ## 角色
                        您是“南航”航空公司的客户聊天支持代理。请以友好、乐于助人且愉快的方式来回复。
                        您正在通过在线聊天系统与客户互动。
                        今天的日期是 {current_date}。
                        
                        ## 要求
                        1. 在提供有关预订航班或取消预订的信息之前，您必须始终从用户处获取以下信息：预订号和姓名。
                        2. 请讲中文。
                        
                        """)
                .defaultTools(mcpTools, bookingTools)
                .defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory),
                        new LoggingAdvisor())
                .build();
        log.info("DashScopeChatModel: {}", dashScopeChatModel.getDefaultOptions().getModel());
        log.info("OllamaChatModel: {}", ollamaChatModel.getDefaultOptions().getModel());
        this.ollamaChatModel = ChatClient.builder(ollamaChatModel)
                // .defaultTools(mcpTools)
                .build();
    }


    @GetMapping(value = "/generateStreamAsString", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> generateStreamAsString(@RequestParam(value = "message", defaultValue = "讲个笑话") String message) {

        Flux<String> content = this.dashScopeChatModel.prompt()
                .user(message)
                .system(s -> s.param("current_date", LocalDate.now().toString()))
                .stream()
                .content();
        return content.concatWith(Flux.just("[complete]"));

    }

    @GetMapping(value = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat(@RequestParam(value = "message", defaultValue = "讲个笑话") String message) {

        Flux<String> content = this.ollamaChatModel.prompt()
                .user(message)
                .stream()
                .content();
        return content;
    }

    @GetMapping(value = "/chat2", produces = MediaType.TEXT_HTML_VALUE)
    public String chat2(@RequestParam(value = "message", defaultValue = "讲个笑话") String message) {
        String content = this.ollamaChatModel.prompt()
                .user(message)
                .call()
                .content();
        return content;
    }
}
